package cn.ai.shoes.client.utils;

import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import org.jetbrains.annotations.NotNull;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import java.util.Map;
/**
 * 负载均衡策略
 */
@Component
public class LoadBalanceStrategy {
    @Value("${comfyui.server.strategy:notSwitch}")
    private String strategy;
    @Value("${comfyui.server.num:0}")
    private Integer num;

    public static final String POLLING = "polling";
    public static final String GPU = "GPU";

    @Resource
    private ComfyuiServerUtil comfyuiServerUtil;

    public String switchServer() {
        Map<String, String> serverUrls = comfyuiServerUtil.getAllServerUrls();
        if (POLLING.equals(strategy)) {
            return switchByPolling(serverUrls);
        }else if (GPU.equals(strategy)) {
            return switchByGPU(serverUrls);
        }else {
            return notSwitch(serverUrls,num);
        }
    }
    private String notSwitch(Map<String, String> serverUrls,Integer number) {return LoadBalanceUtils.notSwitch(serverUrls,number);}
    @NotNull
    private String switchByPolling(Map<String, String> serverUrls) {
        return LoadBalanceUtils.switchByPolling(serverUrls);
    }

    private String switchByGPU(Map<String, String> serverUrls) {
        return LoadBalanceUtils.switchByGPU(serverUrls);
    }

    public Map<String, String> getServerUrls() {
        return comfyuiServerUtil.getAllServerUrls();
    }
}
